package com.saascloud.datasource.core;

import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import com.alibaba.druid.pool.DruidDataSource;
import com.saascloud.core.saas.SaasContextUtil;
import lombok.AllArgsConstructor;
import lombok.Data;
import lombok.NoArgsConstructor;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.core.env.Environment;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;
import org.springframework.util.StringUtils;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.HashMap;
import java.util.Map;
import java.util.Objects;

/**
 * @Author: liangxu
 * @Contact: one_dream@aliyun.com
 * @Date: 2022/5/17 14:40
 * @Version: 1.0
 * @Description:
 */
public class TenantDataSource extends AbstractRoutingDataSource {
    private static Logger logger = LoggerFactory.getLogger(TenantDataSource.class);
    private Environment environment;

    public TenantDataSource(Environment environment) {
        this.environment = environment;
    }

    private Map<Object, Object> tenantDataSources = new HashMap<>(16);

    @Override
    protected Object determineCurrentLookupKey() {
        String tenantId = SaasContextUtil.getTenantId();
        logger.debug("获取数据源：{}", tenantId);
        return tenantId;
    }

    @Override
    public void afterPropertiesSet() {
        super.setTargetDataSources(buildDataSource(null));
        super.afterPropertiesSet();
    }

    @Override
    public DataSource determineTargetDataSource() {
        Object lookupKey = this.determineCurrentLookupKey();
        DataSource dataSource = (DataSource) tenantDataSources.get(lookupKey);
        //无租户信息，走默认数据库配置
        if (dataSource == null && (lookupKey == null)) {
            dataSource = getResolvedDefaultDataSource();
        } else {
            //从数据库查询再次初始化
            buildDataSource(lookupKey);
            dataSource = (DataSource) tenantDataSources.get(lookupKey);
        }
        if (dataSource == null) {
            throw new IllegalStateException("Cannot determine target DataSource for lookup key [" + lookupKey + "]");
        } else {
            return dataSource;
        }
    }

    /**
     * 根据配置文件的配置构建数据源
     *
     * @return
     */
    public Map<Object, Object> buildDataSource(Object dbKey) {
        synchronized (this) {
            if (Objects.nonNull(dbKey)) {
                Object db = tenantDataSources.get(dbKey);
                if (Objects.nonNull(db)) {
                    return tenantDataSources;
                }
            }
            DataSource defaultDataSource = getResolvedDefaultDataSource();
            if (Objects.isNull(defaultDataSource)) {
                String url = environment.getProperty("spring.datasource.url");
                String username = environment.getProperty("spring.datasource.username");
                String password = environment.getProperty("spring.datasource.password");
                DruidDataSource duridDefaultDataSource = new DruidDataSource();
                duridDefaultDataSource.setUrl(url);
                duridDefaultDataSource.setUsername(username);
                duridDefaultDataSource.setPassword(password);
                setDefaultTargetDataSource(duridDefaultDataSource);
                defaultDataSource = duridDefaultDataSource;
            }
            ResultSet resultSet = null;
            PreparedStatement preparedStatement = null;
            Connection connection = null;
            try {
                //查询租户数据库配置
                connection = defaultDataSource.getConnection();
                String sql = "select tenant_id,config_content from tenant_config where config_item='db'";
                if (!StringUtils.isEmpty(dbKey)) {
                    sql = sql + " and tenant_id='" + dbKey + "'";
                }
                preparedStatement = connection.prepareStatement(sql);
                resultSet = preparedStatement.executeQuery();
                while (resultSet.next()) {
                    String tenantId = resultSet.getString("tenant_id");
                    String config = resultSet.getString("config_content");
                    JSONObject jsonConfig = JSONUtil.parseObj(config);
                    DruidDataSource dataSource = new DruidDataSource();
                    dataSource.setUrl(jsonConfig.getStr("url"));
                    dataSource.setUsername(jsonConfig.getStr("username"));
                    dataSource.setPassword(jsonConfig.getStr("password"));
                    tenantDataSources.put(tenantId, dataSource);
                }
                return tenantDataSources;
            } catch (SQLException throwables) {
                throw new RuntimeException(throwables);
            } finally {
                closeQuietly(resultSet);
                closeQuietly(preparedStatement);
                closeQuietly(connection);
            }
        }
    }

    private void closeQuietly(AutoCloseable closeable) {
        if (Objects.nonNull(closeable)) {
            try {
                closeable.close();
            } catch (Exception e) {
                //ignore
            }
        }
    }


    @Data
    @AllArgsConstructor
    @NoArgsConstructor
    static class TenantDbConfig {
        private String url;
        private String username;
        private String password;
    }
}
